package cas.ihep.spark.loop;

import cas.ihep.spark.io.ByteArrayOutputStreamExt;
import cas.ihep.spark.io.ByteBufferUtil;
import cas.ihep.spark.io.Packet;
import cas.ihep.spark.util.LogEnv;
import com.google.common.io.Closer;

import java.io.Closeable;
import java.io.IOException;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.concurrent.ConcurrentHashMap;

@SuppressWarnings("UnstableApiUsage")
public abstract class IterationNetClient implements Closeable{
    Closer closer;

    IterationNetClient(){
        closer=Closer.create();
    }
    @Override
    public void close() throws IOException {
        closer.close();
    }

    public interface VirtualNetContext{
        void write(Packet packet, SocketAddress to)throws IOException;
    }

    public interface OnPacketRead{
        boolean onRead(VirtualNetContext ctx,Packet packet,SocketAddress from);
    }

    public abstract void waitForPacket(OnPacketRead onRead,int timeOut);


    private static class UDPNet extends IterationNetClient implements VirtualNetContext {
        DatagramChannel udp;
        Selector selector;
        ConcurrentHashMap<Integer,Packet> objCache;
        public void close()throws IOException{
            objCache.clear();
            closer.close();
        }

        UDPNet(SocketAddress addr)throws IOException{
            udp=closer.register(DatagramChannel.open());
            udp.bind(addr);
            udp.configureBlocking(false);
            selector=closer.register(Selector.open());
            udp.register(selector, SelectionKey.OP_READ);
            objCache=new ConcurrentHashMap<>();
        }

        @Override
        public void waitForPacket(OnPacketRead onRead, int timeOut) {
            boolean isRunning=true;
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            try {
                while (isRunning && selector.select(timeOut) > 0) {
                    Iterator<SelectionKey> iter=selector.selectedKeys().iterator();
                    while(iter.hasNext()) {
                        SelectionKey key = iter.next();
                        iter.remove();
                        if (key.isReadable()) {
                            SocketAddress from=((DatagramChannel)key.channel()).receive(buffer);
                            buffer.flip();
                            Packet packet=Packet.wrap(buffer);
                            int tp=packet.type();
                            if(tp==Packet.DataRequest){
                                int iternum=packet.numOfIterations();
                                Packet result=objCache.get(iternum);
                                if(result==null) {
                                    isRunning = onRead.onRead(this, packet, from);
                                }else{
                                    ByteBuffer ibuffer=result.getBuffer();
                                    udp.send(ibuffer,from);
                                    ibuffer.rewind();
                                }
                            }else{
                                isRunning=onRead.onRead(this,packet,from);
                            }
                            buffer.clear();
                        }
                    }
                }
            }catch (IOException e){
                e.printStackTrace();
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        @Override
        public void write(Packet packet, SocketAddress to) throws IOException {
            if(packet.type()==Packet.DataReply) {
                objCache.put(packet.numOfIterations(), packet);
            }
            ByteBuffer buf=packet.getBuffer();
            udp.send(buf,to);
        }
    }

    private static class TCPNet extends IterationNetClient implements VirtualNetContext{
        SocketChannel tcp;
        SocketAddress from;
        Selector selector;

        TCPNet(SocketAddress addr)throws IOException{
            tcp=closer.register(SocketChannel.open(from=addr));
            tcp.configureBlocking(false);
            selector=closer.register(Selector.open());
            tcp.register(selector,SelectionKey.OP_READ);
            LogEnv.getInstance().debug("Connecting to "+addr.toString());
        }

        @Override
        public void waitForPacket(OnPacketRead onRead, int timeOut) {
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            try{
                ByteArrayOutputStreamExt cache=new ByteArrayOutputStreamExt(32);
                boolean running=true;
                ByteBuffer tmpBuf=null;
                while(running){
                    if(selector.select(timeOut)>0){
                        Iterator<SelectionKey> iter=selector.selectedKeys().iterator();
                        while(iter.hasNext()){
                            SelectionKey key=iter.next();
                            iter.remove();
                            if(key.isReadable()){
                                SocketChannel chn= (SocketChannel) key.channel();
                                chn.read(buffer);
                                buffer.flip();
                                cache.write(buffer);
                                buffer.clear();
                            }
                        }
                    }
                    while(cache.size()>0){
                        int bodyLength=cache.preadInt(0);
                        int totalLength=bodyLength+4;
                        if(cache.size()>=totalLength){
                            if(tmpBuf==null || tmpBuf.capacity()<totalLength){
                                tmpBuf=ByteBuffer.allocate(totalLength);
                            }
                            tmpBuf.limit(totalLength);
                            cache.readAndRemove(tmpBuf);
                            tmpBuf.flip();
                            if(!onRead.onRead(this,Packet.wrap(tmpBuf),from)){
                                running=false;
                                break;
                            }
                            tmpBuf.clear();
                            /*byte[] buf=new byte[bodyLength+4];
                            assert(cache.readAndRemove(buf)==buf.length);
                            if(!onRead.onRead(thisContext,new Packet(buf),from)){
                                running=false;
                                break;
                            }*/
                        }else{
                            break;
                        }
                    }
                }

            }catch (IOException e){
                e.printStackTrace();
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        @Override
        public void write(Packet packet, SocketAddress to) throws IOException {
            tcp.write(packet.getBuffer());
        }
    }

    public static IterationNetClient newClient(String type, SocketAddress local, SocketAddress remote)throws IOException{
        if("udp".equalsIgnoreCase(type)){
            LogEnv.getInstance().debug("Creating udp client(%s)",local);
            return new UDPNet(local);
        }else if("tcp".equalsIgnoreCase(type)){
            LogEnv.getInstance().debug("Creating tcp client(%s)",remote);
            return new TCPNet(remote);
        }
        throw new IOException("Unsupported type of client: "+type);
    }
}
